/*
 *    Copyright 2006-2022 the original author or authors.
 *
 *    Licensed under the Apache License, Version 2.0 (the "License");
 *    you may not use this file except in compliance with the License.
 *    You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *    Unless required by applicable law or agreed to in writing, software
 *    distributed under the License is distributed on an "AS IS" BASIS,
 *    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *    See the License for the specific language governing permissions and
 *    limitations under the License.
 */
package org.mybatis.generator.api;

import static org.mybatis.generator.internal.util.ClassloaderUtility.getCustomClassloader;
import static org.mybatis.generator.internal.util.messages.Messages.getString;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.nio.charset.Charset;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.mybatis.generator.codegen.RootClassInfo;
import org.mybatis.generator.config.Configuration;
import org.mybatis.generator.config.Context;
import org.mybatis.generator.config.MergeConstants;
import org.mybatis.generator.exception.InvalidConfigurationException;
import org.mybatis.generator.exception.ShellException;
import org.mybatis.generator.internal.DefaultShellCallback;
import org.mybatis.generator.internal.ObjectFactory;
import org.mybatis.generator.internal.XmlFileMergerJaxp;

/**
 * This class is the main interface to MyBatis generator. A typical execution of
 * the tool involves these steps:
 *
 * <ol>
 * <li>Create a Configuration object. The Configuration can be the result of a
 * parsing the XML configuration file, or it can be created solely in Java.</li>
 * <li>Create a MyBatisGenerator object</li>
 * <li>Call one of the generate() methods</li>
 * </ol>
 *
 * @author Jeff Butler
 * @see org.mybatis.generator.config.xml.ConfigurationParser
 */
public class MyBatisGenerator {

	private static final ProgressCallback NULL_PROGRESS_CALLBACK = new ProgressCallback() {
	};

	private final Configuration configuration;

	private final ShellCallback shellCallback;

	private final List<GeneratedJavaFile> generatedJavaFiles = new ArrayList<>();

	private final List<GeneratedXmlFile> generatedXmlFiles = new ArrayList<>();

	private final List<GeneratedKotlinFile> generatedKotlinFiles = new ArrayList<>();

	/**
	 * Any kind of generated file generated by plugin methods
	 * contextGenerateAdditionalFiles.
	 */
	private final List<GeneratedFile> otherGeneratedFiles = new ArrayList<>();

	private final List<String> warnings;

	private final Set<String> projects = new HashSet<>();

	/**
	 * Constructs a MyBatisGenerator object.
	 *
	 * @param configuration The configuration for this invocation
	 * @param shellCallback an instance of a ShellCallback interface. You may
	 *                      specify <code>null</code> in which case the
	 *                      DefaultShellCallback will be used.
	 * @param warnings      Any warnings generated during execution will be added to
	 *                      this list. Warnings do not affect the running of the
	 *                      tool, but they may affect the results. A typical warning
	 *                      is an unsupported data type. In that case, the column
	 *                      will be ignored and generation will continue. You may
	 *                      specify <code>null</code> if you do not want warnings
	 *                      returned.
	 * @throws InvalidConfigurationException if the specified configuration is
	 *                                       invalid
	 */
	public MyBatisGenerator(Configuration configuration, ShellCallback shellCallback, List<String> warnings)
			throws InvalidConfigurationException {
		super();
		if (configuration == null) {
			throw new IllegalArgumentException(getString("RuntimeError.2")); //$NON-NLS-1$
		} else {
			this.configuration = configuration;
		}

		if (shellCallback == null) {
			this.shellCallback = new DefaultShellCallback(false);
		} else {
			this.shellCallback = shellCallback;
		}

		if (warnings == null) {
			this.warnings = new ArrayList<>();
		} else {
			this.warnings = warnings;
		}

		this.configuration.validate();
	}

	/**
	 * This is the main method for generating code. This method is long running, but
	 * progress can be provided and the method can be canceled through the
	 * ProgressCallback interface. This version of the method runs all configured
	 * contexts.
	 *
	 * @param callback an instance of the ProgressCallback interface, or
	 *                 <code>null</code> if you do not require progress information
	 * @throws SQLException         the SQL exception
	 * @throws IOException          Signals that an I/O exception has occurred.
	 * @throws InterruptedException if the method is canceled through the
	 *                              ProgressCallback
	 */
	public void generate(ProgressCallback callback) throws SQLException, IOException, InterruptedException {
		generate(callback, null, null, true);
	}

	/**
	 * This is the main method for generating code. This method is long running, but
	 * progress can be provided and the method can be canceled through the
	 * ProgressCallback interface.
	 *
	 * @param callback   an instance of the ProgressCallback interface, or
	 *                   <code>null</code> if you do not require progress
	 *                   information
	 * @param contextIds a set of Strings containing context ids to run. Only the
	 *                   contexts with an id specified in this list will be run. If
	 *                   the list is null or empty, than all contexts are run.
	 * @throws SQLException         the SQL exception
	 * @throws IOException          Signals that an I/O exception has occurred.
	 * @throws InterruptedException if the method is canceled through the
	 *                              ProgressCallback
	 */
	public void generate(ProgressCallback callback, Set<String> contextIds)
			throws SQLException, IOException, InterruptedException {
		generate(callback, contextIds, null, true);
	}

	/**
	 * This is the main method for generating code. This method is long running, but
	 * progress can be provided and the method can be cancelled through the
	 * ProgressCallback interface.
	 *
	 * @param callback                 an instance of the ProgressCallback
	 *                                 interface, or <code>null</code> if you do not
	 *                                 require progress information
	 * @param contextIds               a set of Strings containing context ids to
	 *                                 run. Only the contexts with an id specified
	 *                                 in this list will be run. If the list is null
	 *                                 or empty, than all contexts are run.
	 * @param fullyQualifiedTableNames a set of table names to generate. The
	 *                                 elements of the set must be Strings that
	 *                                 exactly match what's specified in the
	 *                                 configuration. For example, if table name =
	 *                                 "foo" and schema = "bar", then the fully
	 *                                 qualified table name is "foo.bar". If the Set
	 *                                 is null or empty, then all tables in the
	 *                                 configuration will be used for code
	 *                                 generation.
	 * @throws SQLException         the SQL exception
	 * @throws IOException          Signals that an I/O exception has occurred.
	 * @throws InterruptedException if the method is canceled through the
	 *                              ProgressCallback
	 */
	public void generate(ProgressCallback callback, Set<String> contextIds, Set<String> fullyQualifiedTableNames)
			throws SQLException, IOException, InterruptedException {
		generate(callback, contextIds, fullyQualifiedTableNames, true);
	}

	/**
	 * This is the main method for generating code. This method is long running, but
	 * progress can be provided and the method can be cancelled through the
	 * ProgressCallback interface.
	 *
	 * @param callback                 an instance of the ProgressCallback
	 *                                 interface, or <code>null</code> if you do not
	 *                                 require progress information
	 * @param contextIds               a set of Strings containing context ids to
	 *                                 run. Only the contexts with an id specified
	 *                                 in this list will be run. If the list is null
	 *                                 or empty, than all contexts are run.
	 * @param fullyQualifiedTableNames a set of table names to generate. The
	 *                                 elements of the set must be Strings that
	 *                                 exactly match what's specified in the
	 *                                 configuration. For example, if table name =
	 *                                 "foo" and schema = "bar", then the fully
	 *                                 qualified table name is "foo.bar". If the Set
	 *                                 is null or empty, then all tables in the
	 *                                 configuration will be used for code
	 *                                 generation.
	 * @param writeFiles               if true, then the generated files will be
	 *                                 written to disk. If false, then the generator
	 *                                 runs but nothing is written
	 * @throws SQLException         the SQL exception
	 * @throws IOException          Signals that an I/O exception has occurred.
	 * @throws InterruptedException if the method is canceled through the
	 *                              ProgressCallback
	 */
	public void generate(ProgressCallback callback, Set<String> contextIds, Set<String> fullyQualifiedTableNames,
			boolean writeFiles) throws SQLException, IOException, InterruptedException {

		if (callback == null) {
			callback = NULL_PROGRESS_CALLBACK;
		}

		generatedJavaFiles.clear();
		generatedXmlFiles.clear();
		ObjectFactory.reset();
		RootClassInfo.reset();

		// calculate the contexts to run
		List<Context> contextsToRun;
		if (contextIds == null || contextIds.isEmpty()) {
			contextsToRun = configuration.getContexts();
		} else {
			contextsToRun = new ArrayList<>();
			for (Context context : configuration.getContexts()) {
				if (contextIds.contains(context.getId())) {
					contextsToRun.add(context);
				}
			}
		}

		// setup custom classloader if required
		if (!configuration.getClassPathEntries().isEmpty()) {
			ClassLoader classLoader = getCustomClassloader(configuration.getClassPathEntries());
			ObjectFactory.addExternalClassLoader(classLoader);
		}

		// now run the introspections...
		int totalSteps = 0;
		for (Context context : contextsToRun) {
			totalSteps += context.getIntrospectionSteps();
		}
		callback.introspectionStarted(totalSteps);

		for (Context context : contextsToRun) {
			context.introspectTables(callback, warnings, fullyQualifiedTableNames);
		}

		// now run the generates
		totalSteps = 0;
		for (Context context : contextsToRun) {
			totalSteps += context.getGenerationSteps();
		}
		callback.generationStarted(totalSteps);

		for (Context context : contextsToRun) {
			context.generateFiles(callback, generatedJavaFiles, generatedXmlFiles, generatedKotlinFiles,
					otherGeneratedFiles, warnings);
		}

		// now save the files
		if (writeFiles) {
			callback.saveStarted(generatedXmlFiles.size() + generatedJavaFiles.size());

			for (GeneratedXmlFile gxf : generatedXmlFiles) {
				projects.add(gxf.getTargetProject());
				writeGeneratedXmlFile(gxf, callback);
			}

			for (GeneratedJavaFile gjf : generatedJavaFiles) {
				projects.add(gjf.getTargetProject());
				writeGeneratedJavaFile(gjf, callback);
			}

			for (GeneratedKotlinFile gkf : generatedKotlinFiles) {
				projects.add(gkf.getTargetProject());
				writeGeneratedFile(gkf, callback);
			}

			for (GeneratedFile gf : otherGeneratedFiles) {
				projects.add(gf.getTargetProject());
				writeGeneratedFile(gf, callback);
			}

			for (String project : projects) {
				shellCallback.refreshProject(project);
			}
		}

		callback.done();
	}

	private void writeGeneratedJavaFile(GeneratedJavaFile gjf, ProgressCallback callback)
			throws InterruptedException, IOException {
		File targetFile;
		String source;
		try {
			File directory = shellCallback.getDirectory(gjf.getTargetProject(), gjf.getTargetPackage());
			targetFile = new File(directory, gjf.getFileName());
			if (targetFile.exists()) {
				if (shellCallback.isMergeSupported()) {
					source = shellCallback.mergeJavaFile(gjf.getFormattedContent(), targetFile,
							MergeConstants.getOldElementTags(), gjf.getFileEncoding());
				} else if (shellCallback.isOverwriteEnabled()) {
					source = gjf.getFormattedContent();
					warnings.add(getString("Warning.11", //$NON-NLS-1$
							targetFile.getAbsolutePath()));
				} else {
					source = gjf.getFormattedContent();
					targetFile = getUniqueFileName(directory, gjf.getFileName());
					warnings.add(getString("Warning.2", targetFile.getAbsolutePath())); //$NON-NLS-1$
				}
			} else {
				source = gjf.getFormattedContent();
			}

			callback.checkCancel();
			callback.startTask(getString("Progress.15", targetFile.getName())); //$NON-NLS-1$
			writeFile(targetFile, source, gjf.getFileEncoding());
		} catch (ShellException e) {
			warnings.add(e.getMessage());
		}
	}

	private void writeGeneratedFile(GeneratedFile gf, ProgressCallback callback)
			throws InterruptedException, IOException {
		File targetFile;
		String source;
		try {
			File directory = shellCallback.getDirectory(gf.getTargetProject(), gf.getTargetPackage());
			targetFile = new File(directory, gf.getFileName());
			if (targetFile.exists()) {
				if (shellCallback.isOverwriteEnabled()) {
					source = gf.getFormattedContent();
					warnings.add(getString("Warning.11", //$NON-NLS-1$
							targetFile.getAbsolutePath()));
				} else {
					source = gf.getFormattedContent();
					targetFile = getUniqueFileName(directory, gf.getFileName());
					warnings.add(getString("Warning.2", targetFile.getAbsolutePath())); //$NON-NLS-1$
				}
			} else {
				source = gf.getFormattedContent();
			}

			callback.checkCancel();
			callback.startTask(getString("Progress.15", targetFile.getName())); //$NON-NLS-1$
			writeFile(targetFile, source, gf.getFileEncoding());
		} catch (ShellException e) {
			warnings.add(e.getMessage());
		}
	}

	private void writeGeneratedXmlFile(GeneratedXmlFile gxf, ProgressCallback callback)
			throws InterruptedException, IOException {
		File targetFile;
		String source;
		try {
			File directory = shellCallback.getDirectory(gxf.getTargetProject(), gxf.getTargetPackage());
			targetFile = new File(directory, gxf.getFileName());
			if (targetFile.exists()) {
				if (gxf.isMergeable()) {
					source = XmlFileMergerJaxp.getMergedSource(gxf, targetFile);
				} else if (shellCallback.isOverwriteEnabled()) {
					source = gxf.getFormattedContent();
					warnings.add(getString("Warning.11", //$NON-NLS-1$
							targetFile.getAbsolutePath()));
				} else {
					source = gxf.getFormattedContent();
					targetFile = getUniqueFileName(directory, gxf.getFileName());
					warnings.add(getString("Warning.2", targetFile.getAbsolutePath())); //$NON-NLS-1$
				}
			} else {
				source = gxf.getFormattedContent();
			}

			callback.checkCancel();
			callback.startTask(getString("Progress.15", targetFile.getName())); //$NON-NLS-1$
			writeFile(targetFile, source, gxf.getFileEncoding());
		} catch (ShellException e) {
			warnings.add(e.getMessage());
		}
	}

	/**
	 * Writes, or overwrites, the contents of the specified file.
	 *
	 * @param file         the file
	 * @param content      the content
	 * @param fileEncoding the file encoding
	 * @throws IOException Signals that an I/O exception has occurred.
	 */
	private void writeFile(File file, String content, String fileEncoding) throws IOException {
		try (FileOutputStream fos = new FileOutputStream(file, false)) {
			OutputStreamWriter osw;
			if (fileEncoding == null) {
				osw = new OutputStreamWriter(fos);
			} else {
				osw = new OutputStreamWriter(fos, Charset.forName(fileEncoding));
			}

			try (BufferedWriter bw = new BufferedWriter(osw)) {
				bw.write(content);
			}
		}
	}

	/**
	 * Gets the unique file name.
	 *
	 * @param directory the directory
	 * @param fileName  the file name
	 * @return the unique file name
	 */
	private File getUniqueFileName(File directory, String fileName) {
		File answer = null;

		// try up to 1000 times to generate a unique file name
		StringBuilder sb = new StringBuilder();
		for (int i = 1; i < 1000; i++) {
			sb.setLength(0);
			sb.append(fileName);
			sb.append('.');
			sb.append(i);

			File testFile = new File(directory, sb.toString());
			if (!testFile.exists()) {
				answer = testFile;
				break;
			}
		}

		if (answer == null) {
			throw new RuntimeException(getString("RuntimeError.3", directory.getAbsolutePath())); //$NON-NLS-1$
		}

		return answer;
	}

	/**
	 * Returns the list of generated Java files after a call to one of the generate
	 * methods. This is useful if you prefer to process the generated files yourself
	 * and do not want the generator to write them to disk.
	 *
	 * @return the list of generated Java files
	 */
	public List<GeneratedJavaFile> getGeneratedJavaFiles() {
		return generatedJavaFiles;
	}

	/**
	 * Returns the list of generated Kotlin files after a call to one of the
	 * generate methods. This is useful if you prefer to process the generated files
	 * yourself and do not want the generator to write them to disk.
	 *
	 * @return the list of generated Kotlin files
	 */
	public List<GeneratedKotlinFile> getGeneratedKotlinFiles() {
		return generatedKotlinFiles;
	}

	/**
	 * Returns the list of generated XML files after a call to one of the generate
	 * methods. This is useful if you prefer to process the generated files yourself
	 * and do not want the generator to write them to disk.
	 *
	 * @return the list of generated XML files
	 */
	public List<GeneratedXmlFile> getGeneratedXmlFiles() {
		return generatedXmlFiles;
	}
}
